{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 7-4. Getting Started with Real-World Robots - Train a policy on your data\n",
    "\n",
    ":::{note}\n",
    "\n",
    "This notebook content is derived from lerobot repo's example note \"[7_get_started_with_real_robot.md](https://github.com/huggingface/lerobot/blob/main/examples/7_get_started_with_real_robot.md#3-record-your-dataset-and-visualize-it)\".\n",
    "\n",
    "The notebook is formated in Jupyter notebook style, so that one can just follow this notebook on Jetson (running Jupyer Lab server inside the lerobot container).\n",
    "\n",
    "Please note that this notebook (along with other in the series) may not be the full copy of the original note, as it tries to catpure the essense and important part that involves with Python code executions.<br>\n",
    "For the full documentation and its update, please always refer to the original document (\"[7_get_started_with_real_robot.md](https://github.com/huggingface/lerobot/blob/main/examples/7_get_started_with_real_robot.md)\").\n",
    "\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## a. Use the `train` script\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Using wandb-core as the SDK backend. Please refer to https://wandb.me/wandb-core for more information.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: W&B API key is configured. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m If you're specifying your api key in code, ensure this code is not shared publicly.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Consider setting the WANDB_API_KEY environment variable, or running `wandb login` from the command line.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Appending key for api.wandb.ai to your netrc file: /root/.netrc\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import wandb\n",
    "wandb.login(key='3f5a229bd42d482c4ff0c6af6a59a727c50e497e')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 2024-10-10 06:04:11 ts/train.py:244 {'dataset_repo_id': 'chitoku/koch_green2',\n",
      " 'device': 'cuda',\n",
      " 'env': {'action_dim': 6,\n",
      "         'fps': '${fps}',\n",
      "         'name': 'real_world',\n",
      "         'state_dim': 6,\n",
      "         'task': None},\n",
      " 'eval': {'batch_size': 50, 'n_episodes': 50, 'use_async_envs': False},\n",
      " 'fps': 30,\n",
      " 'override_dataset_stats': {'observation.images.laptop': {'mean': [[[0.485]],\n",
      "                                                                   [[0.456]],\n",
      "                                                                   [[0.406]]],\n",
      "                                                          'std': [[[0.229]],\n",
      "                                                                  [[0.224]],\n",
      "                                                                  [[0.225]]]},\n",
      "                            'observation.images.phone': {'mean': [[[0.485]],\n",
      "                                                                  [[0.456]],\n",
      "                                                                  [[0.406]]],\n",
      "                                                         'std': [[[0.229]],\n",
      "                                                                 [[0.224]],\n",
      "                                                                 [[0.225]]]}},\n",
      " 'policy': {'chunk_size': 100,\n",
      "            'dim_feedforward': 3200,\n",
      "            'dim_model': 512,\n",
      "            'dropout': 0.1,\n",
      "            'feedforward_activation': 'relu',\n",
      "            'input_normalization_modes': {'observation.images.laptop': 'mean_std',\n",
      "                                          'observation.images.phone': 'mean_std',\n",
      "                                          'observation.state': 'mean_std'},\n",
      "            'input_shapes': {'observation.images.laptop': [3, 480, 640],\n",
      "                             'observation.images.phone': [3, 480, 640],\n",
      "                             'observation.state': ['${env.state_dim}']},\n",
      "            'kl_weight': 10.0,\n",
      "            'latent_dim': 32,\n",
      "            'n_action_steps': 100,\n",
      "            'n_decoder_layers': 1,\n",
      "            'n_encoder_layers': 4,\n",
      "            'n_heads': 8,\n",
      "            'n_obs_steps': 1,\n",
      "            'n_vae_encoder_layers': 4,\n",
      "            'name': 'act',\n",
      "            'output_normalization_modes': {'action': 'mean_std'},\n",
      "            'output_shapes': {'action': ['${env.action_dim}']},\n",
      "            'pre_norm': False,\n",
      "            'pretrained_backbone_weights': 'ResNet18_Weights.IMAGENET1K_V1',\n",
      "            'replace_final_stride_with_dilation': False,\n",
      "            'temporal_ensemble_momentum': None,\n",
      "            'use_vae': True,\n",
      "            'vision_backbone': 'resnet18'},\n",
      " 'resume': False,\n",
      " 'seed': 1000,\n",
      " 'training': {'batch_size': 8,\n",
      "              'delta_timestamps': {'action': '[i / ${fps} for i in '\n",
      "                                             'range(${policy.chunk_size})]'},\n",
      "              'do_online_rollout_async': False,\n",
      "              'eval_freq': -1,\n",
      "              'grad_clip_norm': 10,\n",
      "              'image_transforms': {'brightness': {'min_max': [0.8, 1.2],\n",
      "                                                  'weight': 1},\n",
      "                                   'contrast': {'min_max': [0.8, 1.2],\n",
      "                                                'weight': 1},\n",
      "                                   'enable': False,\n",
      "                                   'hue': {'min_max': [-0.05, 0.05],\n",
      "                                           'weight': 1},\n",
      "                                   'max_num_transforms': 3,\n",
      "                                   'random_order': False,\n",
      "                                   'saturation': {'min_max': [0.5, 1.5],\n",
      "                                                  'weight': 1},\n",
      "                                   'sharpness': {'min_max': [0.8, 1.2],\n",
      "                                                 'weight': 1}},\n",
      "              'log_freq': 100,\n",
      "              'lr': 1e-05,\n",
      "              'lr_backbone': 1e-05,\n",
      "              'num_workers': 4,\n",
      "              'offline_steps': 80000,\n",
      "              'online_buffer_capacity': None,\n",
      "              'online_buffer_seed_size': 0,\n",
      "              'online_env_seed': None,\n",
      "              'online_rollout_batch_size': 1,\n",
      "              'online_rollout_n_episodes': 1,\n",
      "              'online_sampling_ratio': 0.5,\n",
      "              'online_steps': 0,\n",
      "              'online_steps_between_rollouts': 1,\n",
      "              'save_checkpoint': True,\n",
      "              'save_freq': 10000,\n",
      "              'weight_decay': 0.0001},\n",
      " 'use_amp': False,\n",
      " 'video_backend': 'pyav',\n",
      " 'wandb': {'disable_artifact': False,\n",
      "           'enable': True,\n",
      "           'notes': '',\n",
      "           'project': 'lerobot'}}\n",
      "ERROR 2024-10-10 06:04:12 ib/gitlib.py:92 git root error: Cmd('git') failed due to: exit code(128)\n",
      "  cmdline: git rev-parse --show-toplevel\n",
      "  stderr: 'fatal: detected dubious ownership in repository at '/opt/lerobot'\n",
      "To add an exception for this directory, call:\n",
      "\n",
      "\tgit config --global --add safe.directory /opt/lerobot'\n",
      "ERROR 2024-10-10 06:04:12 ib/gitlib.py:92 git root error: Cmd('git') failed due to: exit code(128)\n",
      "  cmdline: git rev-parse --show-toplevel\n",
      "  stderr: 'fatal: detected dubious ownership in repository at '/opt/lerobot'\n",
      "To add an exception for this directory, call:\n",
      "\n",
      "\tgit config --global --add safe.directory /opt/lerobot'\n",
      "\u001b[1m\u001b[34mLogs will be synced with wandb.\u001b[0m\n",
      "INFO 2024-10-10 06:04:13 n/logger.py:132 Track this run --> \u001b[1m\u001b[33mhttps://wandb.ai/chitoku/lerobot/runs/ya962loi\u001b[0m\n",
      "INFO 2024-10-10 06:04:13 ts/train.py:312 make_dataset\n",
      "WARNING 2024-10-10 06:04:13 s/factory.py:68 There might be a mismatch between your training dataset (dataset_repo_id='chitoku/koch_green2') and your environment (cfg.env.name='real_world').\n",
      "INFO 2024-10-10 06:04:13 ts/train.py:328 make_policy\n",
      "WARNING 2024-10-10 06:04:13 s/factory.py:28 Hydra config is missing arguments: {'temporal_ensemble_coeff'}\n",
      "/opt/lerobot/lerobot/scripts/train.py:338: FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. Please use `torch.amp.GradScaler('cuda', args...)` instead.\n",
      "  grad_scaler = GradScaler(enabled=cfg.use_amp)\n",
      "INFO 2024-10-10 06:04:16 on/logger.py:39 \u001b[1m\u001b[33mOutput dir:\u001b[0m outputs/train/act_koch_test\n",
      "INFO 2024-10-10 06:04:16 ts/train.py:349 cfg.env.task=None\n",
      "INFO 2024-10-10 06:04:16 ts/train.py:350 cfg.training.offline_steps=80000 (80K)\n",
      "INFO 2024-10-10 06:04:16 ts/train.py:351 cfg.training.online_steps=0\n",
      "INFO 2024-10-10 06:04:16 ts/train.py:352 offline_dataset.num_samples=10946 (11K)\n",
      "INFO 2024-10-10 06:04:16 ts/train.py:353 offline_dataset.num_episodes=50\n",
      "INFO 2024-10-10 06:04:16 ts/train.py:354 num_learnable_params=51597190 (52M)\n",
      "INFO 2024-10-10 06:04:16 ts/train.py:355 num_total_params=51597238 (52M)\n",
      "INFO 2024-10-10 06:04:16 ts/train.py:421 Start offline training on a fixed dataset\n",
      "INFO 2024-10-10 06:04:43 ts/train.py:192 step:0 smpl:8 ep:0 epch:0.00 loss:81.171 grdn:1118.114 lr:1.0e-05 updt_s:24.799 data_s:2.181\n"
     ]
    }
   ],
   "source": [
    "!cd /opt/lerobot && DATA_DIR=data python lerobot/scripts/train.py \\\n",
    "  dataset_repo_id=chitoku/koch_green2 \\\n",
    "  policy=act_koch_real \\\n",
    "  env=koch_real \\\n",
    "  hydra.run.dir=outputs/train/act_koch_test \\\n",
    "  hydra.job.name=act_koch_test \\\n",
    "  device=cuda \\\n",
    "  wandb.enable=true"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## b. (Optional) Upload policy checkpoints to the hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
